package http.secure.common.keys;

import http.secure.common.constant.DisguiseEnum;
import http.secure.common.constant.RequestTemplate;
import http.secure.common.constant.SecureConstant;
import http.secure.common.constant.SecureProperties;
import http.secure.common.event.entity.KeyEvent;
import http.secure.common.exception.InvalidRequestException;
import http.secure.common.util.Assert;
import http.secure.common.util.SMHelper;
import http.secure.common.util.steganography.BinarySteganography;
import http.secure.common.util.steganography.FontSteganography;
import http.secure.common.util.steganography.ImageSteganography;
import jakarta.servlet.http.HttpServletRequest;
import lombok.RequiredArgsConstructor;
import org.dromara.hutool.http.server.servlet.ServletUtil;
import org.springframework.context.ApplicationEventPublisher;
import org.springframework.stereotype.Component;

import java.util.concurrent.TimeUnit;

/**
 * 密钥交换管理器
 */
@Component
@RequiredArgsConstructor
public class KeySwitcher {
    private static final long DEFAULT_SAVE_TIME = TimeUnit.DAYS.toMillis(1);
    private final ApplicationEventPublisher publisher;
    private final SecureGuidSession guidSession;
    private final SecureProperties properties;
    private final HttpServletRequest request;
    private final SecureKeyStore keyStore;

    public byte[] disguisePubkey(DisguiseEnum disguise, byte[] bytes) {
        // 缓存用户IP和请求时间
        long exchangeMaxTime = properties.getKeyExchangeMaxTime();
        keyStore.save(SecureConstant.IP_ADDR, ServletUtil.getClientIP(request), DEFAULT_SAVE_TIME);
        // 生成SM2密钥对
        String[] keyPair = SMHelper.genSm2Keys();
        String pubkey = keyPair[0], prikey = keyPair[1];
        // 使用对应加密形式
        byte[] data = switch (disguise) {
            case PNG -> ImageSteganography.hideMessage(bytes, pubkey);
            case TTF -> FontSteganography.hideMessage(bytes, pubkey);
            case BINARY -> BinarySteganography.hideMessage(bytes, pubkey);
        };
        // 保存私钥
        String signature = SMHelper.sm3Hash(pubkey);
        keyStore.save(SecureConstant.PRIVATE_KEY + signature, prikey, exchangeMaxTime);
        return data;
    }

    public String exchange(String cpk, String sDigest, String cDigest, String browserSign) {
        // 校验IP一致性
        String savedIp = keyStore.get(SecureConstant.IP_ADDR);
        boolean checkIp = savedIp.equals(ServletUtil.getClientIP(request));
        Assert.isTrue(checkIp, RequestTemplate.IP_ADDRESS_MISMATCH, InvalidRequestException::new);
        // 拼接私钥key
        String priScope = SecureConstant.PRIVATE_KEY + sDigest;
        // 检查交换超时时间
        Long expireTime = keyStore.getExpireTime(priScope);
        long maxTime = properties.getKeyExchangeMaxTime();
        boolean checkTime = expireTime != null && expireTime <= maxTime;
        Assert.isTrue(checkTime, RequestTemplate.MAX_EXCHANGE_TIME, InvalidRequestException::new);
        String prikey = keyStore.get(priScope);
        keyStore.delete(priScope);
        Assert.notNull(prikey, RequestTemplate.PRIVATE_KEY_NOT_FOUND, InvalidRequestException::new);
        // 解密客户端公钥
        String cpubkey = SMHelper.sm2Decrypt(cpk, prikey);
        // 验证服务器签名
        boolean equals = SMHelper.sm3Hash(cpubkey).equals(cDigest);
        Assert.isTrue(equals, RequestTemplate.PUBLIC_KEY_SIGNATURE_MISMATCH, InvalidRequestException::new);
        // 记录客户端公钥和指纹
        keyStore.save(SecureConstant.BROWSER_IV, browserSign, DEFAULT_SAVE_TIME);
        keyStore.save(SecureConstant.CLIENT_PUB_KEY, cpubkey, DEFAULT_SAVE_TIME);
        // 返回加密的SM4密钥
        return rotationKeys(cpubkey);
    }

    /**
     * 更新SM4密钥
     */
    public String rotationKeys(String cpubkey) {
        String sm4key = null;
        if (!properties.isSameSessionExchangeDiffKey()) {
            sm4key = keyStore.get(SecureConstant.SM4_KEY);
        }
        if (sm4key == null) {
            sm4key = SMHelper.genSm4Key();
            keyStore.save(SecureConstant.SM4_KEY, sm4key, properties.getKeyRotationTime() * 1000);
            // 推送密钥交换事件
            String browser = keyStore.get(SecureConstant.BROWSER_IV);
            publisher.publishEvent(new KeyEvent(publisher, guidSession.getGuid(), sm4key, browser));
        }
        return SMHelper.sm2Encrypt(sm4key, cpubkey);
    }
}
